/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#ifndef __AICPU_ZERO_COPY_EXCHANGER_H__
#define __AICPU_ZERO_COPY_EXCHANGER_H__

#include <atomic>
#include <set>
#include <unordered_set>
#include <functional>
#include <hccl/hccl_types.h>
#include "aicpu_operator_pub.h"
#include "coll_alg_param.h"
#include "transport_device_p2p_pub.h"
#include "zero_copy/zero_copy_address_mgr.h"

namespace hccl {
class AicpuZeroCopyExchanger {
public:
    AicpuZeroCopyExchanger(u32 rank, u32 rankSize, const HcclOpResParam *resParam, std::function<bool()> needStop, u32 timeoutSec = 120);
    ~AicpuZeroCopyExchanger();

    HcclResult ExchangeAddress(const std::string &tag, void *localInput, void *localOutput, AlgResourceResponse *algResResponse);

    // yxg-debug 提供提前退出的方法退出阻塞接口
private:
    // 这里使用volatile修饰，是为了避免数据写到CPU cache中，必须要写到内存中
    struct FlagData {
        volatile u64 inAddr;    // 通信的input address
        volatile u64 outAddr;   // 通信的output address
        volatile u64 flag;      // flag标志，表明其那面数据是否有效，必须放置到最后一个成员，这样能保证sdma copy先写数据，后修改flag
    };

    struct TagRes {
        std::set<u32> remoteRanks;    // 本tag所有的通信对端rank
        std::vector<LINK> links;                // 本tag所有的通信对端link
        std::vector<void *> remotePtrs;         // batchSdma copy使用的入参
        std::vector<void *> selfPtrs;
        std::vector<FlagData> selfData;
        std::vector<size_t> sizes;
        std::vector<u32> rankIds;               // 维测信息使用
    };

    static constexpr u64 INVALID_DATA = 0;  // 数据无效
    static constexpr u64 VALID_DATA   = 1;  // 数据有效

    void MemFence() const
    {
        /* 内存屏障，即阻止编译器重排变量读写，也阻止CPU重排变量读写 */
        std::atomic_thread_fence(std::memory_order_seq_cst); 
    }

    // 判断是否所有的ipc都是有效的，如果无效则报错
    bool IsAllIpcAddressValid();
    bool IsSupportZeroCopyLinkType(LinkType linkType);

    /* 尝试从data中读数据，只有flag是valid时可以读, 读成功需要置invalid, 如果失败则返回EAGIN表示需要重试 */
    HcclResult TryToRead(FlagData &data, u64 &in, u64 &out);

    HcclResult GetRemoteRanks(TagRes &tagRes, OpCommTransport &opTransportResponse);
    HcclResult PrepareTagRes(const std::string &tag, OpCommTransport &opTransportResponse);
    HcclResult GetRemoteAddr();
    HcclResult BatchSetLocalAddrToRemote(void *in, void *out);
    HcclResult UpdateTransportAddress();
    std::string DumpLinkInfo(std::set<u32> &doneRanks);

    // 存放进行地址交换后的对端input/output
    u64 inAddrs_[MAX_MODULE_DEVICE_NUM]{};
    u64 outAddrs_[MAX_MODULE_DEVICE_NUM]{};
    u32 rankId_{INVALID_VALUE_RANKID};
    u32 rankSize_{INVALID_VALUE_RANKSIZE};

    const HcclOpResParam *resParam_{nullptr};
    std::function<bool()> needStop_{};
    u32 timeoutSec_ = 120;
    std::unordered_map<std::string, TagRes> tagRes_{};
    TagRes *current_{nullptr};      // 表明当前正在使用的tag资源

    static ZeroCopyAddressMgr globalAddrMgr_;
};
}

#endif